use tokio::task;
use reqwest::Client;
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
use flume::{Sender, Receiver, bounded};
use futures::stream::{FuturesUnordered, StreamExt};
use std::time::Instant;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use std::collections::HashMap;

mod config;

async fn generate_requests(tx: Sender<reqwest::Request>, client: Arc<Client>, id: usize, config: &config::RequestConfig) {
    let mut headers = HashMap::new();
    headers.insert("Content-Type".to_string(), "application/json".to_string());
    headers.insert("User-Agent".to_string(), "MyApp/1.0".to_string());

    let header_map = map_to_header_map(headers);
    let request = client
            .request(
                reqwest::Method::from_bytes(config.method.as_bytes()).unwrap(),
                &config.address
            )
            .version(reqwest::Version::HTTP_11)
            .headers(header_map)
            .body(config.body.clone().unwrap_or_default())
            .build()
            .unwrap();
    loop {
        let request = request.try_clone().unwrap();
        if tx.send_async(request).await.is_err() {
            break;
        }
        println!("Producer {} sent a request", id);
    }
}
fn map_to_header_map(headers: HashMap<String, String>) -> HeaderMap {
    let mut header_map = HeaderMap::new();
    for (key, value) in headers {
        // 尝试从字符串创建 HeaderName
        if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
            // 尝试从字符串创建 HeaderValue
            if let Ok(header_value) = HeaderValue::from_str(&value) {
                // 插入头部到 HeaderMap 中
                header_map.insert(header_name, header_value);
            } else {
                // 处理 HeaderValue 创建失败的情况（值无效）
                println!("无效的头部值: {}", value);
            }
        } else {
            // 处理 HeaderName 创建失败的情况（无效的头部名称）
            println!("无效的头部名称: {}", key);
        }
    }
    header_map
}
async fn send_requests_for_duration(
    rx: Receiver<reqwest::Request>,
    id: usize,
    total_counter: Arc<AtomicUsize>,
    success_counter: Arc<AtomicUsize>,
    client: Arc<Client>
) {
    while let Ok(req) = rx.recv_async().await {
        let client = client.clone();
        let total_counter = total_counter.clone();
        let success_counter = success_counter.clone();
        tokio::spawn(async move {
            total_counter.fetch_add(1, Ordering::SeqCst);
            let response = client.execute(req).await;
            if response.is_ok() {
                success_counter.fetch_add(1, Ordering::SeqCst);
            }
            println!("Consumer {} sent a request", id);
        });
    }
}

async fn monitor_requests(
    total_counter: Arc<AtomicUsize>,
    success_counter: Arc<AtomicUsize>,
    rx: Receiver<reqwest::Request>
) {
    let mut last_total = 0;
    let start_time = Instant::now();
    loop {
        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
        let total = total_counter.load(Ordering::SeqCst);
        let success = success_counter.load(Ordering::SeqCst);
        let elapsed_secs = start_time.elapsed().as_secs();
        let success_rate = if total > 0 {
            success as f64 / total as f64 * 100.0
        } else {
            0.0
        };
        let per_sec = total - last_total;
        last_total = total;
        let remaining = rx.len();
        println!(
            "Total requests: {}, Success: {}, Success rate: {:.2}%, Requests per second: {}, Remaining in channel: {}",
            total, success, success_rate, per_sec, remaining
        );
    }
}

async fn start_producers(
    tx: Sender<reqwest::Request>,
    client: Arc<Client>,
    config: &config::RequestConfig
) -> FuturesUnordered<task::JoinHandle<()>> {
    let mut producers = FuturesUnordered::new();
    for i in 0..config.threads {
        let tx = tx.clone();
        let client = Arc::clone(&client);
        let config = config.clone();
        producers.push(task::spawn(async move {
            generate_requests(tx, client, i, &config).await;
        }));
    }
    producers
}

async fn start_consumers(
    rx: Receiver<reqwest::Request>,
    total_counter: Arc<AtomicUsize>,
    success_counter: Arc<AtomicUsize>,
    client: Arc<Client>,
    config: &config::RequestConfig
) -> FuturesUnordered<task::JoinHandle<()>> {
    let mut consumers = FuturesUnordered::new();
    for i in 0..config.threads {
        let rx = rx.clone();
        let total_counter = Arc::clone(&total_counter);
        let success_counter = Arc::clone(&success_counter);
        let client = Arc::clone(&client);
        consumers.push(task::spawn(async move {
            send_requests_for_duration(rx, i, total_counter, success_counter, client).await;
        }));
    }
    consumers
}

#[tokio::main(flavor = "multi_thread", worker_threads = 4)]
async fn main() {
    let client = Arc::new(Client::new());
    let (tx, rx) = bounded(1000);
    let total_counter = Arc::new(AtomicUsize::new(0));
    let success_counter = Arc::new(AtomicUsize::new(0));

    // 选择一个请求配置，可以根据需求动态选择
    let request_config = &config::CONFIG.requests_config[0]; // 这里使用第一个请求配置

    // 启动生产者和消费者
    let producers = start_producers(tx.clone(), Arc::clone(&client), request_config).await;
    let consumers = start_consumers(rx.clone(), Arc::clone(&total_counter), Arc::clone(&success_counter), Arc::clone(&client), request_config).await;

    // 启动监控任务
    if config::CONFIG.is_monitor {
        let monitor_handle = task::spawn(async move {
            monitor_requests(total_counter, success_counter, rx).await;
        });

        // 等待一段时间后关闭生产者和消费者
        tokio::time::sleep(tokio::time::Duration::from_secs(request_config.duration_seconds)).await;

        // 关闭发送端以停止消费者
        drop(tx);

        // 等待所有任务完成
        producers.for_each(|_| async {}).await;
        consumers.for_each(|_| async {}).await;

        // 监控任务不需要等待，因为它是一个无限循环
        // 在真实场景中，你可能需要一个机制来停止监控任务
        let _ = monitor_handle.await;
    }
}
